# coding=utf-8
# 代码文件：chapter9/ch9.2.6.py
import os

from PyPDF2 import PdfFileReader, PdfFileWriter

# 设置输入目录
indir = r'data\in'
#  设置输出目录
outdir = r'data\out'


# 查找dir目录下ext后缀名的文件列表
# dir参数是文件所在目录，ext参数是文件后缀名

def findext(dir, ext):
    allfile = os.listdir(dir)
    # 返回过滤器对象
    files_filter = filter(lambda x: x.endswith(ext), allfile)
    # 从过滤器对象提取列表
    list2 = list(files_filter)
    return list2  # 返回过滤后条件文件名


# 添加水印函数
# namelist要添加水印的文件名，output参数是输出文件名
def add_watermark(namelist):
    # 水印文件
    watermark = r'data/画画程序员水印.pdf'
    # 创建读取水印文件PdfFileReader对象
    watermarkpdf_reader = PdfFileReader(watermark)
    # 读取水印文件第一个页面对象
    watermark_page = watermarkpdf_reader.getPage(0)

    # 遍历每一个添加水印的PDF文件
    for name in namelist:
        infile = os.path.join(indir, name)

        outfile = os.path.join(outdir, name)
        pdf_reader = PdfFileReader(infile)
        pdf_writer = PdfFileWriter()

        # 遍历PDF文档每一个页面
        for page_no in range(pdf_reader.getNumPages()):
            # 获得输入的PDF文档一个页面对象
            page = pdf_reader.getPage(page_no)
            # 给页面添加水印
            page.mergePage(watermark_page)

            # 将数据写入到PDF文档
            pdf_writer.addPage(page)

        # 写入文件
        with open(outfile, 'wb') as out:
            pdf_writer.write(out)


if __name__ == '__main__':
    # 查找indir目录中所有PDF文件
    list2 = findext(indir, '.pdf')
    add_watermark(list2)
    print('添加水印完成！')
